#include "MDX.h"

#include <iomanip>
#include <iostream>
#include <sstream>

#include "utf8.h"

// 构造函数
MDX::MDX(const std::string& fname, const std::string& encoding, bool substyle,
         const std::pair<std::vector<uint8_t>, std::string>& passcode)
    : MDictBase(fname, encoding, passcode) {
    m_substyle = substyle;
    cout << m_substyle << " qss " << m_stylesheet.empty() << endl;
    parseRecordBlockInfo();
}

// 处理记录数据
std::vector<uint8_t> MDX::_treat_record_data(const std::vector<uint8_t>& data) {
    if (data.empty()) {
        return {};
    }

    // 对数据进行样式替换
    std::vector<uint8_t> treated_data = data;

    if (m_substyle && !m_stylesheet.empty()) {
        treated_data = _substitute_stylesheet(data);
    }

    return treated_data;
}

// 替换样式表
std::vector<uint8_t> MDX::_substitute_stylesheet(
    const std::vector<uint8_t>& txt) {
    std::string text(txt.begin(), txt.end());

    // 查找<head>标签并插入样式表
    size_t head_pos = text.find("<head>");
    if (head_pos != std::string::npos) {
        head_pos += 6;  // 跳过<head>标签

        // 构建样式表内容
        std::string stylesheet_content = "<style type='text/css'>\n";
        for (const auto& style : m_stylesheet) {
            stylesheet_content += style.second.second + "\n";
        }
        stylesheet_content += "</style>\n";

        // 插入样式表
        text.insert(head_pos, stylesheet_content);
    } else {
        // 如果没有<head>标签，添加完整的HTML结构
        std::string stylesheet_content = "<style type='text/css'>\n";
        for (const auto& style : m_stylesheet) {
            stylesheet_content += style.second.second + "\n";
        }
        stylesheet_content += "</style>\n";

        text = "<html><head>" + stylesheet_content + "</head><body>" + text +
               "</body></html>";
    }

    // 转换回字节数组
    return std::vector<uint8_t>(text.begin(), text.end());
}

void MDX::parseRecordBlockInfo() {
    if (m_version >= 3) {
    } else {
        std::ifstream file(m_filename, std::ios::binary);
        if (!file) {
            throw std::runtime_error("Failed to open file: " + m_filename);
        }

        file.seekg(m_record_block_offset);

        // 读取记录块信息
        int num_record_blocks = (m_version >= 2.0)
                                    ? BinaryIO::read_value<uint64_t>(file)
                                    : BinaryIO::read_value<uint32_t>(file);
        int num_entries = (m_version >= 2.0)
                              ? BinaryIO::read_value<uint64_t>(file)
                              : BinaryIO::read_value<uint32_t>(file);

        if (num_entries != m_num_entries) {
            throw std::runtime_error("Number of entries mismatch");
        }
        cout << "record blocks的数量: " << num_record_blocks << " "
             << num_entries << endl;

        int record_block_info_size = (m_version >= 2.0)
                                         ? BinaryIO::read_value<uint64_t>(file)
                                         : BinaryIO::read_value<uint32_t>(file);
        int record_block_size = (m_version >= 2.0)
                                    ? BinaryIO::read_value<uint64_t>(file)
                                    : BinaryIO::read_value<uint32_t>(file);

        // 读取记录块信息列表
        std::vector<std::pair<uint64_t, uint64_t>> record_block_info_list;
        int size_counter = 0;
        uint64_t offset = 0, tellg_ = 0;

        for (int i = 0; i < num_record_blocks; ++i) {
            uint64_t compressed_size =
                (m_version >= 2.0) ? BinaryIO::read_value<uint64_t>(file)
                                   : BinaryIO::read_value<uint32_t>(file);
            uint64_t decompressed_size =
                (m_version >= 2.0) ? BinaryIO::read_value<uint64_t>(file)
                                   : BinaryIO::read_value<uint32_t>(file);
            record_block_info_list.emplace_back(compressed_size,
                                                decompressed_size);
            size_counter += m_number_width * 2;
            m_recordInfos.push_back(
                {tellg_, offset, decompressed_size, compressed_size});
            offset += decompressed_size;
            tellg_ += compressed_size;
        }

        if (size_counter != record_block_info_size) {
            throw std::runtime_error("Record block info size mismatch");
        }
        m_record_chunks_offset = file.tellg();
    }
}

// 读取记录
// record block header
// record block info
// record block 1
// record block 2
// ...
std::vector<std::pair<std::string, std::vector<uint8_t>>> MDX::_read_records() {
    std::vector<std::pair<std::string, std::vector<uint8_t>>> result;

    if (m_version >= 3) {
        // v3版本实现（简化版）
        std::ifstream file(m_filename, std::ios::binary);
        if (!file) {
            throw std::runtime_error("Failed to open file: " + m_filename);
        }

        // 读取记录数据
        file.seekg(m_record_block_offset);
        uint32_t number = BinaryIO::read_value<uint32_t>(file);
        uint64_t total_size = BinaryIO::read_value<uint64_t>(file);

        // 处理每个记录块
        for (uint32_t i = 0; i < number; ++i) {
            uint32_t decompressed_size = BinaryIO::read_value<uint32_t>(file);
            uint32_t compressed_size = BinaryIO::read_value<uint32_t>(file);

            std::vector<uint8_t> block_data =
                BinaryIO::read_bytes(file, compressed_size);
            std::vector<uint8_t> decompressed_block_data =
                _decode_block(block_data, decompressed_size);

            // 简化实现：直接将整个块作为一个记录
            if (i < m_key_list.size()) {
                const std::string& key_text = m_key_list[i].second;
                std::vector<uint8_t> treated_data =
                    _treat_record_data(decompressed_block_data);
                result.emplace_back(key_text, treated_data);
            }
        }

        file.close();
    } else {
        // v1/v2版本实现
        std::ifstream file(m_filename, std::ios::binary);
        if (!file) {
            throw std::runtime_error("Failed to open file: " + m_filename);
        }

        file.seekg(m_record_block_offset);

        // 读取记录块信息
        int num_record_blocks = (m_version >= 2.0)
                                    ? BinaryIO::read_value<uint64_t>(file)
                                    : BinaryIO::read_value<uint32_t>(file);
        int num_entries = (m_version >= 2.0)
                              ? BinaryIO::read_value<uint64_t>(file)
                              : BinaryIO::read_value<uint32_t>(file);

        if (num_entries != m_num_entries) {
            throw std::runtime_error("Number of entries mismatch");
        }
        cout << "record blocks的数量: " << num_record_blocks << " "
             << num_entries << endl;

        int record_block_info_size = (m_version >= 2.0)
                                         ? BinaryIO::read_value<uint64_t>(file)
                                         : BinaryIO::read_value<uint32_t>(file);
        int record_block_size = (m_version >= 2.0)
                                    ? BinaryIO::read_value<uint64_t>(file)
                                    : BinaryIO::read_value<uint32_t>(file);

        // 读取记录块信息列表
        std::vector<std::pair<uint64_t, uint64_t>> record_block_info_list;
        int size_counter = 0;

        for (int i = 0; i < num_record_blocks; ++i) {
            uint64_t compressed_size =
                (m_version >= 2.0) ? BinaryIO::read_value<uint64_t>(file)
                                   : BinaryIO::read_value<uint32_t>(file);
            uint64_t decompressed_size =
                (m_version >= 2.0) ? BinaryIO::read_value<uint64_t>(file)
                                   : BinaryIO::read_value<uint32_t>(file);
            record_block_info_list.emplace_back(compressed_size,
                                                decompressed_size);
            size_counter += m_number_width * 2;
        }

        if (size_counter != record_block_info_size) {
            throw std::runtime_error("Record block info size mismatch");
        }

        // 处理记录块
        long offset = 0;
        int i = 0;
        size_counter = 0;
        int idx = -1;

        for (const auto& block_info : record_block_info_list) {
            idx++;
            uint64_t compressed_size = block_info.first;
            uint64_t decompressed_size = block_info.second;

            // 1. 根据record block info提供的每个记录块的长度信息，
            // 从文件中分割提取出记录块,并解压缩
            vector<uint8_t> record_block_data =
                BinaryIO::read_bytes(file, compressed_size);
            vector<uint8_t> record_block = _decode_block(
                record_block_data, static_cast<uint32_t>(decompressed_size));

            // cout << "记录块" << i << "字节数:" << compressed_size << endl;
            // 根据键块中的偏移信息分割记录块
            while (i < m_key_list.size()) {
                uint64_t record_start = m_key_list[i].first;
                const std::string& key_text = m_key_list[i].second;
                // if (key_text == "output") {
                //     cout << "记录块(output)" << ": " << idx << endl;
                // }

                // 检查是否到达当前记录块的末尾
                if (record_start - offset >= record_block.size()) {
                    break;
                }

                // 确定记录结束位置
                uint64_t record_end;
                if (i < m_key_list.size() - 1) {
                    record_end = m_key_list[i + 1].first;
                } else {
                    record_end = record_block.size() + offset;
                }
                // 提取记录数据
                std::vector<uint8_t> data(
                    record_block.begin() + (record_start - offset),
                    record_block.begin() + (record_end - offset));

                // 处理记录数据
                std::vector<uint8_t> treated_data = _treat_record_data(data);

                // 添加到结果中
                result.emplace_back(key_text, treated_data);

                i++;
            }

            offset += record_block.size();
            size_counter += compressed_size;
        }

        if (size_counter != record_block_size) {
            throw std::runtime_error("Record block size mismatch");
        }

        file.close();
    }

    return result;
}

// 获取所有键值对
std::vector<std::pair<std::string, std::vector<uint8_t>>> MDX::items() {
    return _read_records();
}

string MDX::query(string key) {
    string record;
    // 键块解析出的pair<uint64_t, string>为一个单词，其中第一个数值为偏移量
    // 这里的偏移量是指的是全部(record block)解压之后的偏移量

    // 快速查询,减少内存占用的算法：
    // 1. 根据key block获取 key list列表: m_key_list
    // 2. 根据key list中查询到某个key对应的offset偏移量
    // 3. 读取record block info信息，获取所有记录块的解压缩后的大小
    // 4. 计算offset在哪一个记录块中
    // 5. 解压对应的记录块并根据下一个key对应的offset确定记录的结束位置

    uint64_t offset = 0, size = 0;
    // if (auto it = m_key_offset.find(key); it != m_key_offset.end()) {
    //     offset = it->second.offset;
    //     size = it->second.size;
    // } else {
    //     cout << "not found key" << endl;
    //     return "";
    // }

    Record r = get_offset_and_size(key);
    if (r.size == 0 && r.offset == 0) {
        cout << "not found key" << endl;
        return "";
    }
    offset = r.offset;
    size = r.size;

    // 查找key对应的offset所在的记录块
    int i = 0;
    for (auto info : m_recordInfos) {
        if (info.offset_start > offset) {
            break;
        }
        i++;
    }
    // cout << "find key(output) : " << i << endl;
    RecordInfo info = m_recordInfos[i - 1];

    // 直接读取查找到的块
    std::ifstream file(m_filename, std::ios::binary);
    if (!file) {
        throw std::runtime_error("Failed to open file: " + m_filename);
    }

    file.seekg(m_record_chunks_offset + info.tellg_);
    vector<uint8_t> record_block_data =
        BinaryIO::read_bytes(file, info.compressed_size);
    file.close();
    vector<uint8_t> record_block = _decode_block(
        record_block_data, static_cast<uint32_t>(info.decompressed_size));

    // 定位到键对应的记录
    vector<uint8_t> data(
        record_block.begin() + (offset - info.offset_start),
        record_block.begin() + (offset - info.offset_start + size));
    record = string(data.begin(), data.begin() + data.size() - 1);
    // cout << record << endl;

    if (encoding() == "GBK") {
        record = EncodeCvt::Convert(record.data(), "gbk", "utf8");
    }
    return record;
}

string MDX::metaInfo() {
    auto header = _read_header();
    std::ostringstream oss;
    for (auto& i : header) {
        oss << i.first << ": " << i.second << "\n";
    }
    return oss.str();
}

void MDX::test() {
    /* 1. 硬编码 GBK 字节流  */
    unsigned char gbkBytes[] = {0xd6, 0xd0, 0xce, 0xc4};  // “中文”
    std::string gbk(reinterpret_cast<char*>(gbkBytes), 4);

    /* 2. 打印原始字节 */
    std::cout << "原始 GBK  : ";
    for (unsigned char c : gbk)
        std::cout << std::hex << std::setw(2) << std::setfill('0')
                  << static_cast<int>(c) << ' ';
    std::cout << '\n';

    /* 3. 转 UTF-8 */
    std::string utf8 = EncodeCvt::Convert(gbk.c_str(), "gbk", "utf8");

    /* 4. 打印结果 */
    std::cout << "转换 UTF-8: ";
    for (unsigned char c : utf8)
        std::cout << std::hex << std::setw(2) << std::setfill('0')
                  << static_cast<int>(c) << ' ';
    std::cout << '\n';
}
